Knowledge Bases for Amazon Bedrock (with OpenSearch Serverless)をSAMで実装してみた

つくぼし

2024.05.26

こんにちは、つくぼし(tsukuboshi0755)です！

Knowledge Bases for Amazon Bedrockが東京リージョンにやってきた事で、RAG構成の候補としてKnowledge Basesをより使いやすくなったかと思います。

以前Knowledge Basesを用いたRAG構成として、ベクターストアにAuroraを利用したパターンをCloudFormationで作成しました。

ただ多くの方々は、Knowledge Basesコンソールのクイック作成でも用いられる、ベクターストアにOpenSearch Serverlessを利用したパターンで作成する事が多いのではないでしょうか。

そこで今回は、ベクターストアにOpenSearch Serverlessを利用したパターンをSAMで実装してみたので、紹介したいと思います！

なおインデックスのセットアップの実施で一部カスタムリソースを含む形になるため、その点にご留意頂けるとありがたいです。

構成

今回SAMで作成する構成図は以下のとおりです。

Knowledge Bases for Amazon Bedrockに必要なリソースとして、以下のリソースも合わせて作成します。

ベクターストア：OpenSearch Serverless
データソース：S3

またLambda-Backed Custom Resourceを利用して、OpenSearch Serverlessに対するインデックスのセットアップを自動化します。

テンプレート

全体のコードは、以下のリポジトリに格納しています。

以下では、テンプレートの主要な部分を抜粋して説明します。

OpenSearch Serverless

  OSSCollection:
    Type: 'AWS::OpenSearchServerless::Collection'
    Properties:
      Name: !Sub '${AWS::StackName}-col'
      Type: VECTORSEARCH
      StandbyReplicas: !Ref OSSCollectionStandbyReplicas
    DependsOn: OSSEncryptionPolicy
  OSSEncryptionPolicy:
    Type: 'AWS::OpenSearchServerless::SecurityPolicy'
    Properties:
      Name: !Sub '${AWS::StackName}-ep'
      Type: encryption
      Policy: !Sub >-
        {
            "Rules": [
                {
                    "Resource": [
                        "collection/${AWS::StackName}-col"
                    ],
                    "ResourceType": "collection"
                }
            ],
            "AWSOwnedKey": true
        }
  OSSNetworkPolicy:
    Type: 'AWS::OpenSearchServerless::SecurityPolicy'
    Properties:
      Name: !Sub '${AWS::StackName}-np'
      Type: network
      Policy: !Sub >-
        [
            {
                "Rules": [
                    {
                        "Resource": [
                            "collection/${AWS::StackName}-col"
                        ],
                        "ResourceType": "dashboard"
                    },
                    {
                        "Resource": [
                            "collection/${AWS::StackName}-col"
                        ],
                        "ResourceType": "collection"
                    }
                ],
                "AllowFromPublic": true
            }
        ]
  OSSDataAccessPolicy:
    Type: 'AWS::OpenSearchServerless::AccessPolicy'
    Properties:
      Name: !Sub '${AWS::StackName}-dp'
      Type: data
      Policy: !Sub
        - >-
            [
                {
                    "Rules": [
                        {
                            "Resource": [
                                "collection/${AWS::StackName}-col"
                            ],
                            "Permission": [
                                "aoss:CreateCollectionItems",
                                "aoss:UpdateCollectionItems",
                                "aoss:DescribeCollectionItems"
                            ],
                            "ResourceType": "collection"
                        },
                        {
                            "Resource": [
                                "index/${AWS::StackName}-col/*"
                            ],
                            "Permission": [
                                "aoss:CreateIndex",
                                "aoss:UpdateIndex",
                                "aoss:DescribeIndex",
                                "aoss:ReadDocument",
                                "aoss:WriteDocument"
                            ],
                            "ResourceType": "index"
                        }
                    ],
                    "Principal": [
                        "${BedrockKnowledgeBaseRoleArn}",
                        "${CreateIndexFunctionRoleArn}",
                        "${OSSDataAccessPrincipalArn}"
                    ],
                    "Description": ""
                }
            ]
        - BedrockKnowledgeBaseRoleArn: !GetAtt BedrockKnowledgeBaseRole.Arn
          CreateIndexFunctionRoleArn: !GetAtt CreateIndexFunctionRole.Arn

Knowledge Base for Amazon Bedrockのベクターストアとして必要なOpenSearch Serverlessを作成します。

今回のOpenSearch Serverlessは基本的に以下を参照して作成しているため、詳細についてはこちらのの記事をご参照ください。

なお上記のブログからの変更点として、データアクセスポリシーに対して後述のKnowledge Bases及びインデックスセットアップ関数に対するロールにも、アクセス権限を付与するようにしています。

S3 Bucket

  DataSourceBucket:
    Type: AWS::S3::Bucket
    Properties:
      BucketName: !Sub '${AWS::StackName}-ds-bucket-${AWS::AccountId}'
      PublicAccessBlockConfiguration:
        BlockPublicAcls: true
        BlockPublicPolicy: true
        IgnorePublicAcls: true
        RestrictPublicBuckets: true
      BucketEncryption:
        ServerSideEncryptionConfiguration:
          - ServerSideEncryptionByDefault:
              SSEAlgorithm: AES256
            BucketKeyEnabled: true
      OwnershipControls:
        Rules:
          - ObjectOwnership: BucketOwnerEnforced

Knowledge Base for Amazon Bedrockのデータソースとして必要なS3バケットを作成します。

このバケットに対してユーザーがドキュメントをアップロードする事で、Knowledge Base for Amazon Bedrockがデータソースとして利用できるようになります。

IAM Policy

  BedrockFMPolicyForKB:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      ManagedPolicyName: !Sub ${AWS::StackName}-bedrock-fm-policy
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Sid: BedrockInvokeModelStatement
            Effect: Allow
            Action:
              - 'bedrock:InvokeModel'
            Resource: !Sub arn:aws:bedrock:${AWS::Region}::foundation-model/${EmbeddingModelId}
      Roles:
        - !Ref BedrockKnowledgeBaseRole

  BedrockOSSPolicyForKB:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      ManagedPolicyName: !Sub ${AWS::StackName}-bedrock-oss-policy
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Sid: OpenSearchServerlessAPIAccessAllStatement
            Effect: Allow
            Action:
              - 'aoss:APIAccessAll'
            Resource: !GetAtt OSSCollection.Arn
      Roles:
        - !Ref BedrockKnowledgeBaseRole
        - !Ref CreateIndexFunctionRole

  BedrockS3PolicyForKB:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      ManagedPolicyName: !Sub ${AWS::StackName}-bedrock-s3-policy
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Sid: S3ListBucketStatement
            Effect: Allow
            Action:
              - 's3:ListBucket'
            Resource: !GetAtt  DataSourceBucket.Arn
            Condition:
              StringEquals:
                aws:ResourceAccount: !Ref 'AWS::AccountId'
          - Sid: S3GetObjectStatement
            Effect: Allow
            Action:
              - 's3:GetObject'
            Resource:
              - !Join
                - ''
                - - 'arn:aws:s3:::'
                  - !Ref DataSourceBucket
                  - '/*'
            Condition:
              StringEquals:
                aws:ResourceAccount: !Ref 'AWS::AccountId'
      Roles:
        - !Ref BedrockKnowledgeBaseRole

Lambda-Backed Custom Resource及びKnowledge Base for Amazon Bedrockに必要なIAMポリシーを作成します。

Bedrock、Secrets Manager、OpenSearch Serverless、S3に対するアクセス権限を設定し、IAMロールにアタッチします。

Lambda-Backed Custom Resource

SAMテンプレート

  SetupOSSIndex:
    Type: 'Custom::OpenSearchServerlessIndex'
    Properties:
      ServiceToken: !GetAtt CreateIndexFunction.Arn
      Region: !Ref 'AWS::Region'
      Dimension: !If [IsTitanEmbedTextV1, 1536, 1024]
      EmbeddingModelId: !Ref EmbeddingModelId
      CollectionId: !GetAtt OSSCollection.Id
      IndexName: !FindInMap [IndexMap, OSSIndexName, Name]
      VectorField: !FindInMap [IndexMap, OSSVectorField, Name]
      MappingText: !FindInMap [IndexMap, OSSMappingText, Name]
      MappingMetadata: !FindInMap [IndexMap, OSSMappingMetadata, Name]
    DependsOn:
      - BedrockOSSPolicyForKB

  CreateIndexFunctionRole:
    Type: 'AWS::IAM::Role'
    Properties:
      RoleName: !Sub ${AWS::StackName}-create-index-function-role
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: 'sts:AssumeRole'
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole

  CreateIndexFunction:
    Type: AWS::Serverless::Function
    Properties:
      FunctionName: !Sub '${AWS::StackName}-create-index-function'
      Handler: index.lambda_handler
      Role: !GetAtt CreateIndexFunctionRole.Arn
      Runtime: python3.12
      Timeout: 600
      LoggingConfig:
        LogFormat: JSON
        ApplicationLogLevel: INFO
        SystemLogLevel: INFO
      Layers:
        - !Ref OSSLayer
      CodeUri: function/

  OSSLayer:
    Type: AWS::Serverless::LayerVersion
    Properties:
      LayerName: !Sub '${AWS::StackName}-create-index-function-layer'
      ContentUri: layer/
      CompatibleRuntimes:
        - python3.12
    Metadata:
      BuildMethod: python3.12

Lambdaコード

import logging
import time
from typing import Any, Dict

import boto3
import cfnresponse
from opensearchpy import AWSV4SignerAuth, OpenSearch, RequestsHttpConnection

logger = logging.getLogger()

def lambda_handler(event: Dict[str, Any], context: Any) -> None:
    request_type = event["RequestType"]
    logger.info(f"RequestType: {request_type}")

    try:
        if request_type == "Create":
            # Get the parameters from the event
            region = event["ResourceProperties"]["Region"]
            dimension = event["ResourceProperties"]["Dimension"]
            colection_id = event["ResourceProperties"]["CollectionId"]
            index_name = event["ResourceProperties"]["IndexName"]
            vector_field = event["ResourceProperties"]["VectorField"]
            mapping_text = event["ResourceProperties"]["MappingText"]
            mapping_metadata = event["ResourceProperties"]["MappingMetadata"]

            # Create OpenSearch client
            service = "aoss"
            host = colection_id + "." + region + "." + service + "." + "amazonaws.com"
            credentials = boto3.Session().get_credentials()
            auth = AWSV4SignerAuth(credentials, region, service)

            client = OpenSearch(
                hosts=[{"host": host, "port": 443}],
                http_auth=auth,
                use_ssl=True,
                verify_certs=True,
                connection_class=RequestsHttpConnection,
                pool_maxsize=20,
            )

            # Create OpenSearch Index (Wait for 30 seconds for index creation)
            index_body = {
                "settings": {
                    "index": {
                        "number_of_shards": "2",
                        "number_of_replicas": "0",
                        "knn": "true",
                    }
                },
                "mappings": {
                    "properties": {
                        vector_field: {
                            "type": "knn_vector",
                            "dimension": dimension,
                            "method": {
                                "name": "hnsw",
                                "space_type": "l2",
                                "engine": "faiss",
                                "parameters": {},
                            },
                        },
                        mapping_text: {
                            "type": "text",
                        },
                        mapping_metadata: {"type": "text", "index": "false"},
                    }
                },
            }

            response = client.indices.create(
                index_name,
                body=index_body,
            )

            logger.info(f"Index creating: {response}")
            time.sleep(30)

            cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
        if request_type == "Update":
            cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
        if request_type == "Delete":
            cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
    except Exception as e:
        cfnresponse.send(event, context, cfnresponse.FAILED, {"Message": str(e)})

今回のテンプレートの肝となる箇所です。

Knowledge Base for Amazon BedrockのベクターストアとしてOpenSearch Serverlessを利用する際、別途インデックスのセットアップが必要となります。

こちらはCloudFormationでは対応していないため、カスタムリソースを使用してリクエストタイプがCreateとなる時のみ、OpenSearch Serverlessに対してOpenSearch APIをLambda関数で実行する事で、インデックスのセットアップを自動化します。

カスタムリソース及びリクエストタイプについては、以下の記事をご参照ください。

また今回はOpenSearch APIとしてopensearch-pyモジュールを使用しています。

そのためrequirements.txtにopensearch-pyを記載し、Lambda関数のデプロイ時にレイヤーとして追加しています。

opensearch-pyについては、以下の公式ドキュメントをご参照ください。

カスタムリソースとして定義したLambda関数では、以下のセットアップ処理についてOpenSearch APIを実行します。

コレクションへの接続
ベクターストア用のインデックス作成

カスタムリソースによりインデックスが正常に作成されているかどうかを確認するには、OpenSearchコンソールを使用します。

作成したコレクションに対して、以下のようなインデックスが確認できればOKです。

またコンソールだけでなく、awscurlまたはcurlコマンド等を用いる事で、インデックスの詳細設定を確認できます。

CloudShell等の環境で簡単に確認できるため、必要に応じてご参照ください。

# エンドポイントの指定

AOSS_ENDPOINT=<OpenSearch Serverlessのエンドポイント>
INDEX_NAME=bedrock-knowledge-base-default-index

# awscurlのインストール

pip install awscurl

# インデックスの確認

awscurl --service aoss --region ${AWS_REGION} \
  -X GET  \
  ${AOSS_ENDPOINT}/${INDEX_NAME} | jq -r .

上記を実施した際に、以下のようなインデックスの詳細設定が確認できればOKです。

{
  "bedrock-knowledge-base-default-index": {
    "aliases": {},
    "mappings": {
      "properties": {
        "AMAZON_BEDROCK_METADATA": {
          "type": "text",
          "index": false
        },
        "AMAZON_BEDROCK_TEXT_CHUNK": {
          "type": "text"
        },
        "bedrock-knowledge-base-default-vector": {
          "type": "knn_vector",
          "dimension": 1024, // ベクトルの次元数 1536 or 1024
          "method": {
            "engine": "faiss",
            "space_type": "l2",
            "name": "hnsw",
            "parameters": {}
          }
        }
      }
    },
    "settings": {
      "index": {
        "number_of_shards": "2",
        "provided_name": "bedrock-knowledge-base-default-index",
        "knn": "true",
        "creation_date": "XXXXXXXXXXXXX",
        "number_of_replicas": "0",
        "uuid": "xxxxxxxxxxxxxxxxxxxx",
        "version": {
          "created": "XXXXXXXXX"
        }
      }
    }
  }
}

Knowledge Base for Amazon Bedrock

  BedrockKnowledgeBaseRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Sub ${AWS::StackName}-bedrock-kb-role
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: [bedrock.amazonaws.com]
            Action: ['sts:AssumeRole']

  BedrockKnowledgeBase:
    Type: AWS::Bedrock::KnowledgeBase
    Properties:
      Name: !Sub ${AWS::StackName}-knowledge-base
      KnowledgeBaseConfiguration:
          Type: VECTOR
          VectorKnowledgeBaseConfiguration:
            EmbeddingModelArn: !Sub arn:aws:bedrock:${AWS::Region}::foundation-model/${EmbeddingModelId}
      RoleArn: !GetAtt BedrockKnowledgeBaseRole.Arn
      StorageConfiguration:
        Type: OPENSEARCH_SERVERLESS
        OpensearchServerlessConfiguration:
            CollectionArn: !GetAtt OSSCollection.Arn
            FieldMapping:
              MetadataField: !FindInMap [IndexMap, OSSMappingMetadata, Name]
              TextField: !FindInMap [IndexMap, OSSMappingText, Name]
              VectorField: !FindInMap [IndexMap, OSSVectorField, Name]
            VectorIndexName: !FindInMap [IndexMap, OSSIndexName, Name]
    DependsOn: SetupOSSIndex

  BedrockKnowledgeBaseDS:
    Type: AWS::Bedrock::DataSource
    Properties:
      DataDeletionPolicy: RETAIN
      KnowledgeBaseId: !Ref BedrockKnowledgeBase
      Name: !Sub ${AWS::StackName}-data-source
      DataSourceConfiguration:
        Type: S3
        S3Configuration:
          BucketArn: !GetAtt DataSourceBucket.Arn

最後にKnowledge Base for Amazon Bedrockを作成し、事前に作成したOpenSearch Serverless、Secrets Manager、S3を紐づけます。

EmbeddingModeArnでは、Knowledge Base for Amazon Bedrockに対して利用する埋め込みモデルのARNを指定します。

またOpensearchServerlessConfigurationでは、OpenSearch Serverlessに対するコレクション及びインデックス情報を指定します。

動作確認

上記のテンプレートをデプロイする事で、Knowledge Base for Amazon Bedrockを利用して正常に文書データを検索できるか確認します。

なお今回は、最近Knowledge Base for Amazon Bedrockが利用可能になった東京リージョンap-northeast-1でデプロイを行います。

SAM テンプレートのデプロイ

以下のコマンドを実行し、SAMテンプレートをデプロイします。


#　プリンシパルARNを指定
IAM_PRINCIPAL_ARN=<IAMプリンシパルARN>

# SAM テンプレートのデプロイ

sam build
sam deploy --parameter-overrides OSSDataAccessPrincipalArn=${IAM_PRINCIPAL_ARN}

EmbeddingModelIdでは、Knowledge Base for Amazon Bedrockに対して利用する埋め込みモデルのIDを指定します。

今回はデフォルトのcohere.embed-multilingual-v3を指定します。

OSSCollectionStandbyReplicasは、OpenSearch Serverlessのスタンバイレプリカの有効/無効を指定します。

こちらも今回はデフォルトのDISABLEDを指定します。

OSSDataAccessPrincipalArnでは、OpenSearch Serverlessに対してアクセス権限を持つIAMプリンシパル(ユーザー/ロール)のARNを指定します。

こちらは指定必須であるため、OpenSearch Serverlessにアクセスが必要なIAMプリンシパルのARNを指定してください。

なおOpenSearch Serverlessのデプロイ時間が長く、およそ20分程度かかります。

S3バケットへのデータ投入

次にCloudShell上で以下のコマンドを実施し、PDFファイルを作成したS3バケットにアップロードします。

# S3 バケット名を設定

BUCKET_NAME=<バケット名>

# AWS の公式ドキュメントの PDF ファイルをダウンロード

mkdir -p pdf
cd pdf
wget https://docs.aws.amazon.com/ja_jp/amazondynamodb/latest/developerguide/dynamodb-dg.pdf -O DynamoDB.pdf
wget https://docs.aws.amazon.com/ja_jp/lambda/latest/dg/lambda-dg.pdf -O Lambda.pdf
wget https://docs.aws.amazon.com/ja_jp/vpc/latest/userguide/vpc-ug.pdf -O VPC.pdf
wget https://docs.aws.amazon.com/ja_jp/bedrock/latest/userguide/bedrock-ug.pdf -O Bedrock.pdf
cd ..

# S3 バケットに PDF ファイルをアップロード

aws s3 cp pdf s3://${BUCKET_NAME} --recursive

ファイルのアップロードが成功すると、以下の通り作成したS3バケットに4つのPDFファイルが作成されます。